import pandas as pd
from trader.tools.MyTT import MA


K_DAY = "day"
K_WEEK = "week"
K_MONTH = "month"
K_5m = "5m"
K_15m = "15m"
K_30m = "30m"
K_60m = "60m"

K_TYPE = {
    K_DAY: 1,
    K_WEEK: 2,
    K_MONTH: 3,
    K_5m: 5,
    K_15m: 15,
    K_30m: 30,
    K_60m: 60,
}


def normalize_code(stock_code: str) -> str:
    """
    Normalize stock code to the format used by akshare.
    """
    if stock_code.startswith("6"):
        return f"sh{stock_code}"
    if stock_code.find("300", 0, 4) == 0:
        return f"sz{stock_code}"
    if stock_code.find("200", 0, 4) == 0:
        return f"sz{stock_code}"

    return f"sz{stock_code}"

def bs_code(stock_code: str) -> str:
    """
    Normalize stock code to the format used by akshare.
    """
    if stock_code.startswith("6"):
        return f"sh.{stock_code}"
    if stock_code.find("300", 0, 4) == 0:
        return f"sz.{stock_code}"
    if stock_code.find("200", 0, 4) == 0:
        return f"sz.{stock_code}"

    return f"sz.{stock_code}"


def wrap_stat(df, type_="MA") -> pd.DataFrame:
    indicator = {
        "MA": [20, 34, 55, 610],
    }

    indexs = indicator[type_]
    for i in indexs:
        df[f"close_{i}_ma"] = MA(df["close"], i)

    return df
